package com.will.yuaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.will.yuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author will
 * @since 2025/7/24 19:02
 * 处理工具调用的抽象代理类，具体实现了think和act方法，可以用作创建实例的父类
 */
@EqualsAndHashCode(callSuper = true)
@Slf4j
@Data
public class ToolCallAgent extends ReActAgent{

    // 工具调用
    private final ToolCallback[] availableTools;

    //保存工具调用的响应结果
    private ChatResponse toolCallResponse;

    // 工具调用的管理者
    private final ToolCallingManager toolCallManager;

    //禁用SpringAI的默认工具调用机制，使用自定义的
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools){
        super();
        this.availableTools = availableTools;
        this.toolCallManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                //使用自定义工具调用
                .withInternalToolExecutionEnabled(false)
                .build();
    }

    @Override
    public boolean think() {
        //校验提示词，拼接用户提示词
        if (StrUtil.isNotBlank(getNextStepPrompt())){
            Message userMessage = new UserMessage(getNextStepPrompt());
            getMessageList().add(userMessage);
        }
        try {
            //调用大模型
            List<Message> messageList = getMessageList();
            Prompt prompt = new Prompt(messageList, this.chatOptions);
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .toolCallbacks(availableTools)
                    .call()
                    .chatResponse();
            //记录响应，用于等下Act
            this.toolCallResponse = chatResponse;
            //解析工具调用结果
            //助手消息
            AssistantMessage assistantMessage = null;
            if (chatResponse != null) {
                assistantMessage = chatResponse.getResult().getOutput();
            }
            //获取要调用的工具
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            //输出提示消息
            String result = assistantMessage.getText();
            log.info(getName() + "的思考：" + result);
            log.info(getName() + "选择了 " + toolCallList.size() + " 个工具进行调用");
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s,参数：%s", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            if (toolCallList.isEmpty()){
                //只有不调用工具时，才手动添加助手消息
                getMessageList().add(assistantMessage);
                return false;
            }else {
                //需要调用工具时，无需添加助手消息，因为调用工具是会自动记录
                return true;
            }
        } catch (Exception e) {
            log.error("{}的思考过程遇到的问题：{}", getName(), e.getMessage());
            getMessageList().add(new AssistantMessage("思考过程遇到问题:" + e.getMessage()));
            return false;
        }
    }

    @Override
    public String act() {
        if (!toolCallResponse.hasToolCalls()){
            return "result: 没有工具调用结果";
        }
        Prompt prompt = new Prompt(getMessageList(), this.chatOptions);
        //调用工具
        ToolExecutionResult toolExecutionResult = toolCallManager.executeToolCalls(prompt, toolCallResponse);
        //记录消息上下文
        setMessageList(toolExecutionResult.conversationHistory());
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        //判断是否调用了终止工具
        if (toolResponseMessage.getResponses()
                .stream()
                .anyMatch(response -> response.name().equals("doTerminate"))){
            //任务结束，状态改为完成
            setState( AgentState.FINISHED);
        }

        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() + "\nresult: " + response.responseData())
                .collect(Collectors.joining("\n"));
        log.info(results);
        return results;
    }
}
